# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon, inkyu

import asyncio
import logging

import kafka
import kafka.common
from tornado.queues import Queue
from tornado.ioloop import IOLoop

from coin.base.param_util import to_list


class AsyncKafkaPublisher(object):
  def __init__(self, *, kafka_config, error_callback=None, logger=None):
    self._logger = logger or logging.getLogger(__name__)
    self._error_callback = error_callback or (lambda topic, msg: None)

    self._kafka_config = kafka_config
    assert kafka_config is not None

    self._running = False
    self._kafka_producer = None
    self._ioloop = None

  def publish(self, topic, *msgs):
    # TODO(inkyu): Multithread support
    assert self._running
    for msg in msgs:
      assert isinstance(msg, bytes)
      self._queue.put_nowait((topic, msg))

  def start(self, ioloop=None):
    assert not self._running
    self._ioloop = IOLoop.current()
    self._queue = Queue(maxsize=0)
    self._ioloop.add_callback(self._run_loop)
    self._running = True

  async def _init_kafka_producer(self):
    while True:
      try:
        self._kafka_producer = kafka.KafkaProducer(
            bootstrap_servers=self._kafka_config.kafka_servers,
            retries=5,
            batch_size=1024,
            max_request_size=1048576,
            max_in_flight_requests_per_connection=1)
        return
      except Exception as e:
        self._logger.error('Kafka error(%s): Retrying in 5 seconds.', e)
        await asyncio.sleep(5)

  async def _run_loop(self):
    async for (topic, msg) in self._queue:
      if not self._running:
        break

      if self._kafka_producer is None:
        await self._init_kafka_producer()

      self._kafka_send(topic, msg)
      self._queue.task_done()

  def stop(self):
    if not self._running:
      return
    self._running = False
    self._queue.put((None, None))
    self._kafka_clean_up()

  def _kafka_clean_up(self):
    if self._kafka_producer:
      self._kafka_producer.close()
    self._kafka_producer = None

  def _kafka_send_error(self, topic, msg, error, *args, **kwargs):
    self._logger.error('Send error: (%s)', error)
    self._kafka_clean_up()
    self._error_callback(topic, msg)

  def _kafka_send(self, topic, msg, *args, **kwargs):
    try:
      f = self._kafka_producer.send(topic, msg, *args, **kwargs)
      f.add_errback(self._ioloop.add_callback, self._kafka_send_error, topic, msg)
    except kafka.common.KafkaTimeoutError as e:
      logging.exception('Error sending message')
      self._kafka_send_error(topic, msg, e)
